from typing import List, Dict, Iterator, Callable, Union

import heapdict
import pyshark
from pyshark.packet.packet import Packet

from errors import ErrorCounterException
from flow import Flow, get_flow_info, generate_stream_id_by_ip, pop_outdated_streams
from vpn_analyzers import VPNAnalyzer


class TransportAnalyzer:
    """ TransportAnalyzer: Analyzer for analyzing multiple transport layers """

    def __init__(self,
                 udp_timeout: int = 5,
                 gre_timeout: int = 5,
                 esp_timeout: int = 5,
                 l2tp_timeout: int = 5):
        """ an analyzer to generate all kinds of flows in the transport layer

        :param udp_timeout: timeout of udp flow
        :param gre_timeout: timeout of gre flow
        :param esp_timeout: timeout of esp flow
        :param l2tp_timeout: timeout of l2tp flow
        """

        # These are the parameters in the initialization
        self.udp_timeout = udp_timeout
        self.gre_timeout = gre_timeout
        self.esp_timeout = esp_timeout
        self.l2tp_timeout = l2tp_timeout

        # These attributes will be used in the analyzing method
        self.tcp_fin_count = {}  # type: Dict[int, int]
        self.udp_stream_lru_packet_time = heapdict.heapdict()
        self.gre_stream_lru_packet_time = heapdict.heapdict()
        self.esp_stream_lru_packet_time = heapdict.heapdict()
        self.l2tp_stream_lru_packet_time = heapdict.heapdict()

        self.stream_lru_packet_time_dict = {
            "udp": (self.udp_stream_lru_packet_time, self.udp_timeout),
            "gre": (self.gre_stream_lru_packet_time, self.gre_timeout),
            "esp": (self.esp_stream_lru_packet_time, self.esp_timeout),
            "l2tp": (self.l2tp_stream_lru_packet_time, self.l2tp_timeout),
        }

        # These attributes will be used in the stream storage
        self.tcp_streams = {}  # type: Dict[int, List[Packet]]
        self.udp_streams = {}  # type: Dict[int, List[Packet]]
        self.gre_streams = {}  # type: Dict[int, List[Packet]]
        self.esp_streams = {}  # type: Dict[int, List[Packet]]
        self.l2tp_streams = {}  # type: Dict[int, List[Packet]]

        # This is used in the flush method and pop method
        self.stream_dict = {
            "udp": self.udp_streams,
            "gre": self.gre_streams,
            "esp": self.esp_streams,
            "l2tp": self.l2tp_streams,
        }

    def analyze(self, packet: Packet) -> Iterator[Flow]:
        """ main function of TransportAnalyzer. start analyzing procedure

        :param packet: coming packet
        :return: Iterator of flow
        """
        transport_layers = {
            "tcp": self.__tcp_analyze,
            "udp": self.__udp_analyze,
            "gre": self.__gre_analyze,
            "esp": self.__esp_analyze,
            "l2tp": self.__l2tp_analyze,
        }  # type: Dict[str, Callable[[Packet], List[Flow]]]

        for layer_name in transport_layers.keys():
            for layer in packet.layers:
                if layer_name == layer.layer_name:  # if this packet contains one layer defined in the transport_layers
                    layer_flows = transport_layers[layer_name](packet)
                    for layer_flow in layer_flows:
                        if layer_flow.packets is not None and len(layer_flow.packets) > 0:
                            yield layer_flow

    def flush(self) -> Iterator[Flow]:
        """ used when we want to stop the process of TransportAnalyzer. Output all remaining udp/gre/esp flows.

        :return: Iterator of flow
        """
        for flow_type, stream in self.stream_dict.items():
            for stream_index, stream_packets in stream.items():
                yield get_flow_info(stream_packets, stream_index, flow_type)

    def __tcp_analyze(self, packet: Packet) -> List[Flow]:
        """ analyze tcp flow by using (syn,ack) packet and (fin) packet.
        All the three-way handshake and four-way wave-hand packets will be dropped

        :param packet: coming tcp packet
        :return: possible tcp flow. if this packet cannot generate a tcp flow, will return an empty list.
        """

        # Actually, there will only one element in the tcp results
        results = []  # type: List[Flow]

        if int(packet.tcp.flags_syn) == 1 and int(packet.tcp.flags_ack) == 1 and self.tcp_streams.get(
                packet.tcp.stream) is None:
            self.tcp_streams[packet.tcp.stream] = []
        elif int(packet.tcp.flags_fin) == 1 and self.tcp_streams.get(packet.tcp.stream) is not None:
            # if this stream has no tcp fin packet in the past, we will still return None
            if not self.__has_tcp_fin(stream_index=packet.tcp.stream):
                return results

            tcp_stream_packets = self.tcp_streams[packet.tcp.stream]
            del self.tcp_streams[packet.tcp.stream]

            results.append(get_flow_info(tcp_stream_packets, packet.tcp.stream, "tcp"))
        elif self.tcp_streams.get(packet.tcp.stream, None) is not None:
            self.tcp_streams[packet.tcp.stream].append(packet)

        return results

    def __udp_analyze(self, packet: Packet) -> List[Flow]:
        """ analyze udp packets by using udp timeout

        :param packet: coming udp packet
        :return: possible flows. if none of the udp flow is outdated, will return an empty list.
        """
        stream_id = packet.udp.stream

        return self.pop_all_outdated_streams(packet, stream_id, "udp")

    def __gre_analyze(self, packet: Packet) -> List[Flow]:
        """ analyze gre flow by using gre timeout

        :param packet: coming gre packet
        :return: possible flows. if none of the gre flow is outdated, will return an empty list
        """
        # NOTE: in wireshark/tshark, NO concept is called "GRE Stream"
        # However, if we want to have one to analyze, we can only use UDP stream concept to generate a "GRE Stream"
        stream_id = generate_stream_id_by_ip(packet)

        return self.pop_all_outdated_streams(packet, stream_id, "gre")

    def __esp_analyze(self, packet: Packet) -> List[Flow]:
        """ analyze esp flow by using esp timeout

        :param packet: coming esp packet
        :return: possible flows. if none of the esp flow is outdated, will return an empty list.
        """
        # NOTE: In wireshark/tshark, there is no concept called "ESP Stream"
        # However, if we want to have one to analyze, we can only use UDP stream concept to generate an "ESP Stream"
        stream_id = generate_stream_id_by_ip(packet)

        return self.pop_all_outdated_streams(packet, stream_id, "esp")

    def __l2tp_analyze(self, packet: Packet) -> List[Flow]:
        """ analyze l2tp flow by using l2tp timeout

        :param packet: coming l2tp packet
        :return: possible flows. if none of the l2tp flow is outdated, will return an empty list
        """
        # NOTE: In wireshark/tshark, there is no concept called "ESP Stream"
        # However, if we want to have one to analyze, we can only use UDP stream concept to generate an "ESP Stream"
        stream_id = generate_stream_id_by_ip(packet)

        return self.pop_all_outdated_streams(packet, stream_id, "l2tp")

    def __has_tcp_fin(self, stream_index: int) -> bool:
        """ try to detect all the tcp fin status

        :param stream_index: index of the tcp stream in wireshark/tshark
        :return: True if this tcp flow is ended.
        """
        if self.tcp_fin_count.get(stream_index, None) is None:
            self.tcp_fin_count[stream_index] = 1
            return False
        else:
            del self.tcp_fin_count[stream_index]
            return True

    def pop_all_outdated_streams(self, packet, stream_index: int, flow_type: str) -> List[Flow]:
        """ pop out all the outdated udp/gre/esp streams

        :param packet: coming packet
        :param stream_index: index of the stream from the packet layer
        :param flow_type: type name of the flow belonging to the packet.
        :return: possible outdated flows. if none of the flows is outdated, will return an empty list.
        """
        stream = self.stream_dict[flow_type]
        if stream.get(stream_index, None) is None:
            stream[stream_index] = [packet]
        else:
            stream[stream_index].append(packet)

        self.stream_lru_packet_time_dict[flow_type][0][stream_index] = float(packet.sniff_timestamp)

        results = []
        for flow_type, (stream_lru_packet_time, timeout) in self.stream_lru_packet_time_dict.items():
            streams = self.stream_dict[flow_type]
            for outdated_flow in pop_outdated_streams(stream_lru_packet_time, streams, packet, flow_type, timeout):
                results.append(outdated_flow)
        return results


class PySharkAnalyzer:
    """ PySharkAnalyzer: analyzer for analyzing packets using pyshark """

    def __init__(self,
                 interface: Union[str, List[str]] = None,
                 file: str = None,
                 vpn_analyzers: List[VPNAnalyzer] = None):
        """ an analyzer using pyshark.

        if interface is given, will analyze packets in the network card of the interface first.
        :param interface: interface name of the network card. if not given, will listen to all the interfaces.
        :param file: pcap file
        """
        if interface is not None:
            self.capture = pyshark.LiveCapture(interface=interface)
        elif file is not None:
            self.capture = pyshark.FileCapture(input_file=file)
        else:
            print("No input, exited")
            return

        self.transport_analyzer = TransportAnalyzer()
        if vpn_analyzers is not None:
            self.vpn_analyzers = vpn_analyzers
        else:
            self.vpn_analyzers = []

    def analyze_packet(self) -> Iterator[Flow]:
        """ analyzing the packets from the network interface card/pcap file into flows.

        :return: Iterator of multiple types of flow. the info of the flow is in the attribute and udps of the flow.
        """
        if isinstance(self.capture, pyshark.LiveCapture):
            packet_generator = self.capture.sniff_continuously()
        elif isinstance(self.capture, pyshark.FileCapture):
            packet_generator = self.capture
        else:
            return

        try:
            for packet in packet_generator:
                for packet_flow in self.transport_analyzer.analyze(packet):
                    for v in self.vpn_analyzers:
                        v.analyze(packet_flow)
                    yield packet_flow
        except KeyboardInterrupt:
            pass
        except EOFError:
            pass
        except ErrorCounterException as e:
            print(e)
        finally:
            for remained_packet_flow in self.transport_analyzer.flush():
                for v in self.vpn_analyzers:
                    v.analyze(remained_packet_flow)
                yield remained_packet_flow
            packet_generator.close()

    def add_vpn_analyzers(self, *analyzers: VPNAnalyzer):
        """
        add vpn analyzers for flow analysis
        :param analyzers: analyzers to append
        :return: None
        """
        for analyzer in analyzers:
            self.vpn_analyzers.append(analyzer)

    def count_packets(self) -> int:
        if not isinstance(self.capture, pyshark.FileCapture):
            return 0

        cnt = 0
        for _ in self.capture:
            cnt += 1

        return cnt
